//============================================================================
// Name        : cfg.c
// Author      : Yugi
// Version     : 1.0
// Description : 
//============================================================================

#include "cfg.h"
#include "slicer.h"

/* 
* Generate a new control flow gragh.
* 
* RETURN:	new cfg
* NOTES:	Return value should be freed by user
*/
PCFG cfg_new(){
	PCFG pcfg = (PCFG) malloc(sizeof(CFG));
	pcfg->nodes = (PBBNODE) malloc(DEFAULT_CFG_SIZE * sizeof(BBNODE));
	pcfg->head = NULL;
	pcfg->tail = NULL;
	pcfg->size = 0;
	pcfg->multiple = 1;
	pcfg->stack = NULL;
	return pcfg;
}

/* 
* Insert a basic block to the control flow gragh.
* 
* REQUIRE:	pcfg not NULL
*			pc in text seg
* RETURN:	Index of the target basic block, if the basic block
*			exist, return the existing basic block index
* MODIFY:	pcfg
*/
unsigned int cfg_add_bb(void *drcontext, PCFG pcfg, app_pc pc){
	/* Iterator temporary variables */
	unsigned int index;		// index for pcfg->nodes
	unsigned int target;	// target basic block index
	instrlist_t *iterator;	// iterator for pcfg->basic_blocks
	app_pc itr_pc;			// iterator pc
	instr_t *pLast;			// last instruction of a basic block
	instr_t *pTemp;			// temp pointer

	/* Iterate through pcfg->basic_blocks */
	for (index = 0; index<pcfg->size; index++){
		iterator = pcfg->nodes[index].basic_block;
		itr_pc = instr_get_app_pc(instrlist_first(iterator));
		pLast = instrlist_last(iterator);

		/*
		* If any existing basic block begins with pc, then 
		* return its index.
		*/
		if (itr_pc == pc)
			return index;

		/*
		* If any existing basic block contains pc, remove 
		* instructions after pc.
		*/
		if (pc>itr_pc && pc<=instr_get_app_pc(pLast)){
			while (instr_get_app_pc(pLast) >= pc){
				pTemp = pLast;
				pLast = instr_get_prev(pLast);
				instrlist_remove(iterator, pTemp);
			}

			target = cfg_add_bb(drcontext, pcfg, pc);
			cfg_add_edge(pcfg, index, target);
			return target;
		}
	}

	/* Add to pcfg->basic_blocks */
	pcfg->nodes[pcfg->size].basic_block = decode_as_bb(drcontext, pc);

	/* Iterate through pcfg->basic_blocks */
	pLast = instrlist_last(pcfg->nodes[pcfg->size].basic_block);

	for (index = 0; index<pcfg->size; index++){
		iterator = pcfg->nodes[index].basic_block;
		itr_pc = instr_get_app_pc(instrlist_first(iterator));

		/*
		* If the current basic block contains the iterator basic
		* block, remove the instruction in current basic block.
		* after itr_pc.
		*/
		if (pc<itr_pc && itr_pc<=instr_get_app_pc(pLast)){
			while (instr_get_app_pc(pLast) >= itr_pc){
				pTemp = pLast;
				pLast = instr_get_prev(pLast);
				instrlist_remove(pcfg->nodes[pcfg->size].basic_block, pTemp);
			}

			cfg_add_edge(pcfg, pcfg->size, index);
		}
	}

	/* Increase pcfg->size */
	pcfg->size++;

	/* If pcfg->size meets the max size, twice it */
	if (pcfg->size == DEFAULT_CFG_SIZE * pcfg->multiple){
		PBBNODE nodes = pcfg->nodes;

		pcfg->multiple++;
		pcfg->nodes = (PBBNODE) malloc(DEFAULT_CFG_SIZE*pcfg->multiple*sizeof(BBNODE));

		for (index = 0; index<pcfg->size; index++)
			pcfg->nodes[index].basic_block = nodes[index].basic_block;

		free(nodes);
	}

	return pcfg->size-1;
}

/* 
* Insert an edge to the control flow gragh.
* 
* REQUIRE:	pcfg not NULL
*			src & dst distinguishes, less than pcfg->size
* RETURN:	true if edge added
* MODIFY:	pcfg
*/
bool cfg_add_edge(PCFG pcfg, unsigned int src, unsigned int dst){
	/* 
	* Iterate through pcfg->edges, note that pcfg->edges is represented
	* by pcfg->head & pcfg->tail
	*/
	PEDGE edge = pcfg->head;
	
	while (edge){
		if (edge->src==src && edge->dst==dst)
			return false;

		edge = edge->next;
	}
	
	/* Add to pcfg->edges */
	edge = (PEDGE) malloc(sizeof(EDGE));
	edge->src = src;
	edge->dst = dst;
	edge->next = NULL;

	if (pcfg->head){
		pcfg->tail->next = edge;
		pcfg->tail = edge;
	}else
		pcfg->head = pcfg->tail = edge;

	return true;
}

/* 
* Static completion. Add branches to pcfg->basic_blocks.
* 
* REQUIRE:	pcfg not NULL
* MODIFY:	pcfg
*/
void cfg_static_complete(void *drcontext, PCFG pcfg){
	unsigned int index;		// index for pcfg->nodes
	instrlist_t *iterator;	// iterator for pcfg->basic_blocks
	instr_t *pLast;			// Last instruction of basic block
	unsigned int target;	// target basic block index
	
	/* iterate through pcfg->basic_blocks */
	for (index = 0; index<pcfg->size; index++){
		iterator = pcfg->nodes[index].basic_block;
		pLast = instrlist_last(iterator);

		switch (instr_get_opcode(pLast)){
			case OP_jmp:
			case OP_jmp_short:
				/* add branch target */
				target = cfg_add_bb(drcontext, pcfg, instr_get_branch_target_pc(pLast));
				cfg_add_edge(pcfg, index, target);
				break;
			case OP_call:
				/* add next pc */
				target = cfg_add_bb(drcontext, pcfg, instr_get_app_pc(pLast) + instr_length(drcontext, pLast));
				cfg_add_edge(pcfg, index, target);
				break;
			case OP_jmp_ind:
			case OP_jmp_far:
			case OP_jmp_far_ind:
			case OP_call_ind:
			case OP_call_far:
			case OP_call_far_ind:
			case OP_ret:
				break;
			default:
				if (instr_is_cti(pLast)){
				/* add branch target and next pc */
					target = cfg_add_bb(drcontext, pcfg, instr_get_branch_target_pc(pLast));
					cfg_add_edge(pcfg, index, target);
					target = cfg_add_bb(drcontext, pcfg, instr_get_app_pc(pLast) + instr_length(drcontext, pLast));
					cfg_add_edge(pcfg, index, target);
				}
		}
	}
}

/* 
* Print control flow gragh
* 
* REQUIRE:	pcfg not NULL
*			outfile not NULL
* NOTE:		if 'content', the instruction will be printed
*/
void cfg_print(void *drcontext, file_t outfile, PCFG pcfg, bool content){
	PEDGE edge = pcfg->head;// Edge iterator
	unsigned int index;		// index for pcfg->basic_blocks
	instr_t *pInstr;		// Instruction iterator

	dr_fprintf(outfile, "digraph CFG {\n");

	/* Only one node in cfg */
	if (!edge){
		if (pcfg->nodes[0].basic_block)
			dr_fprintf(outfile, "\"0x%.8x\"\n", 
				instr_get_app_pc(instrlist_first(pcfg->nodes[0].basic_block)));
	}else if (content)
		/* Iterate through pcfg->edges */
		while (edge){
			dr_fprintf(outfile, "\"0x%.8x-0x%.8x", 
				instr_get_app_pc(instrlist_first(pcfg->nodes[edge->src].basic_block)), 
				instr_get_app_pc(instrlist_last(pcfg->nodes[edge->src].basic_block)));

			/* Print source instruction content */
			pInstr =  instrlist_first(pcfg->nodes[edge->src].basic_block);
			index = 0;

			while (pInstr){
				if (bm_testbit(pcfg->nodes[edge->src].bitmap, index)){
					dr_fprintf(outfile, "\\n");
					instr_disassemble (drcontext, pInstr, outfile);
				}
	
				pInstr = instr_get_next(pInstr);
				index++;
			}

			dr_fprintf(outfile, "\" -> \"0x%.8x-0x%.8x", 
				instr_get_app_pc(instrlist_first(pcfg->nodes[edge->dst].basic_block)),
				instr_get_app_pc(instrlist_last(pcfg->nodes[edge->dst].basic_block)));

			/* Print destination instruction content */
			pInstr =  instrlist_first(pcfg->nodes[edge->dst].basic_block);
			index = 0;

			while (pInstr){
				if (bm_testbit(pcfg->nodes[edge->dst].bitmap, index)){
					dr_fprintf(outfile, "\\n");
					instr_disassemble (drcontext, pInstr, outfile);
				}
	
				pInstr = instr_get_next(pInstr);
				index++;
			}

			dr_fprintf(outfile, "\"\n");
			edge = edge->next;
		}
	else
		while (edge){
			dr_fprintf(outfile, "\"0x%.8x-0x%.8x\" -> \"0x%.8x-0x%.8x\"\n", 
				instr_get_app_pc(instrlist_first(pcfg->nodes[edge->src].basic_block)), 
				instr_get_app_pc(instrlist_last(pcfg->nodes[edge->src].basic_block)),
				instr_get_app_pc(instrlist_first(pcfg->nodes[edge->dst].basic_block)),
				instr_get_app_pc(instrlist_last(pcfg->nodes[edge->dst].basic_block)));
			edge = edge->next;
		}

	dr_fprintf(outfile, "}\n");
}
